from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
import json
import asyncio
from typing import Dict

app = FastAPI()

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.rooms: Dict[str, set] = {}

    async def connect(self, websocket: WebSocket, client_id: str):
        await websocket.accept()
        self.active_connections[client_id] = websocket
        print(f"客户端 {client_id} 已连接")

    def disconnect(self, client_id: str):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
            # 从所有房间中移除
            for room in self.rooms.values():
                room.discard(client_id)
        print(f"客户端 {client_id} 已断开")

    async def send_personal_message(self, message: str, client_id: str):
        if client_id in self.active_connections:
            await self.active_connections[client_id].send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections.values():
            await connection.send_text(message)

    async def join_room(self, room: str, client_id: str):
        if room not in self.rooms:
            self.rooms[room] = set()
        self.rooms[room].add(client_id)

    async def send_to_room(self, message: str, room: str):
        if room in self.rooms:
            for client_id in self.rooms[room]:
                if client_id in self.active_connections:
                    await self.active_connections[client_id].send_text(message)


manager = ConnectionManager()


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket, client_id)
    try:
        while True:
            data = await websocket.receive_text()
            message_data = json.loads(data)

            # 处理不同类型的消息
            message_type = message_data.get("type")

            if message_type == "join_room":
                room = message_data.get("room")
                await manager.join_room(room, client_id)
                await manager.send_to_room(
                    json.dumps({
                        "type": "system",
                        "message": f"用户 {client_id} 加入了房间",
                        "room": room
                    }), room
                )

            elif message_type == "chat":
                room = message_data.get("room")
                await manager.send_to_room(
                    json.dumps({
                        "type": "chat",
                        "sender": client_id,
                        "message": message_data.get("message"),
                        "timestamp": message_data.get("timestamp")
                    }), room
                )

            elif message_type == "broadcast":
                await manager.broadcast(
                    json.dumps({
                        "type": "broadcast",
                        "sender": client_id,
                        "message": message_data.get("message")
                    })
                )

    except WebSocketDisconnect:
        manager.disconnect(client_id)
        await manager.broadcast(
            json.dumps({
                "type": "system",
                "message": f"用户 {client_id} 离开"
            })
        )


@app.get("/")
async def root():
    return {"message": "WebSocket Server is running"}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)